# Created by Gao Song on 2022-03-29

import pandas as pd
from .rq import rq
from .symbols import lq_to_rq_symbol, symbol_type
from datetime import datetime, timedelta
from .trade_date import get_prev_trading_date, is_trading_day, normalize_datetime
from functools import lru_cache

def get_all_future_contracts():
    return rq.all_instruments(type='Future').order_book_id

def get_underlying_symbol(contract_name):
    return ''.join([i for i in contract_name.split('.')[0] if not i.isdigit()]).upper()

def get_current_price_future(symbols):
    rq_symbols = [lq_to_rq_symbol(s) for s in symbols]
    if len(symbols) == 1:
        return pd.Series({symbols[0]: rq.current_snapshot(rq_symbols).last})
    return pd.Series([x.last for x in rq.current_snapshot(rq_symbols)], index=symbols)

def get_current_pct_future(symbols):
    rq_symbols = [lq_to_rq_symbol(s) for s in symbols]

    now = datetime.now()
    if now.hour > 15 and is_trading_day(now):
        barrier_date = normalize_datetime(now)
    else:
        barrier_date = get_prev_trading_date(now)
    barrier = barrier_date.replace(hour=20)

    if len(symbols) == 1:
        tick = rq.current_snapshot(rq_symbols)
        return pd.Series({symbols[0]: tick.last / tick.prev_close - 1 if not tick.datetime < barrier < now else 0})
    return pd.Series(
        [tick.last / tick.prev_close - 1 if not tick.datetime < barrier < now else 0 for tick in rq.current_snapshot(rq_symbols)],
        index=symbols
    )

@lru_cache(maxsize=128)
def get_contract_multiplier(contract_name):
    symbol = get_underlying_symbol(contract_name)
    d = get_prev_trading_date(datetime.now())
    res = rq.futures.get_contract_multiplier(symbol, start_date=d, end_date=d)['contract_multiplier'].squeeze()
    return res

def future_contract_is_night_trading(symbol):
    return rq.instruments(rq.id_convert(symbol)).trading_hours.startswith('21:')

def get_margin_rate(symbol):
    return rq.instruments(rq.id_convert(symbol)).margin_rate

def contract_expiring(symbol):
    if symbol_type(symbol) == 'future':
        sec_id = symbol.split('.')[0]
        today = datetime.now()
        if int(sec_id[-2:]) - today.month in [1, -11] and datetime.today().day > 20:
            return True
    return False
